import cv2
import os
import numpy as np
from ultralytics import YOLO
import time

class YOLOModel1:
    def __init__(self, weights_path="./models/weights/best(2).pt", device="cuda" if cv2.cuda.getCudaEnabledDeviceCount() > 0 else "cpu"):
        """
        初始化 YOLO 分割模型。
        :param weights_path: 模型权重路径
        :param device: 设备 (cpu 或 cuda)
        """
        self.device = device
        self.model = self.load_model(weights_path)
        self.temp_frames_dir = "static/result"  # 存储逐帧结果的临时目录
        os.makedirs(self.temp_frames_dir, exist_ok=True)

    def load_model(self, weights_path):
        """
        加载 YOLO 模型。
        :param weights_path: 模型权重路径
        :return: 加载的模型
        """
        model = YOLO(weights_path)
        model.to(self.device)
        return model

    def segment(self, input_path, output_path, is_video=False):
        """
        目标检测与分割处理函数。
        :param input_path: 输入文件路径（图片或视频）
        :param output_path: 输出文件路径
        :param is_video: 是否为视频
        :return: 统计数据字典
        """
        if is_video:
            return self.process_video(input_path, output_path)
        else:
            return self.process_image(input_path, output_path)

    def process_image(self, input_path, output_path):
        """
        处理单张图片，执行分割任务。
        :param input_path: 输入图片路径
        :param output_path: 输出图片路径
        :return: 统计数据字典
        """
        # 读取图片
        image = cv2.imread(input_path)
        results = self.model(image)  # 推理

        # 提取分割掩码并绘制
        annotated_image = results[0].plot()

        # 保存结果
        cv2.imwrite(output_path, annotated_image)

        # 获取检测统计数据
        detection_stats = self._extract_stats(results)
        return detection_stats

    def process_video(self, input_path, output_path):
        """
        处理视频。
        :param input_path: 输入视频路径
        :param output_path: 输出视频路径
        :return: 每帧统计数据字典
        """
        # 清空临时帧文件夹
        for file in os.listdir(self.temp_frames_dir):
            os.remove(os.path.join(self.temp_frames_dir, file))

        cap = cv2.VideoCapture(input_path)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

        frame_stats = {}  # 用于保存所有帧的统计数据
        frame_count = 0

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # 目标检测
            results = self.model(frame)
            stats = self._extract_stats(results)

            # 按帧编号保存统计结果
            frame_stats[f"frame_{frame_count:04d}"] = stats

            # 绘制检测框
            annotated_frame = results[0].plot()
            out.write(annotated_frame)

            # 保存每一帧到临时文件夹
            frame_path = os.path.join(self.temp_frames_dir, f"frame_{frame_count:04d}.jpg")
            cv2.imwrite(frame_path, annotated_frame)
            frame_count += 1

        cap.release()
        out.release()

        # 返回每帧的检测统计数据
        return frame_stats


    def _extract_stats(self, results):
        """
        提取检测结果的统计数据。
        :param results: YOLO 检测结果
        :return: 统计数据字典
        """
        labels, counts = np.unique(results[0].boxes.cls.cpu().numpy(), return_counts=True)
        detection_stats = {self.model.names[int(label)]: int(count) for label, count in zip(labels, counts)}
        return detection_stats

    def get_next_frame(self):
        """
        返回下一个检测结果帧的路径和统计数据，用于前端逐帧更新。
        :return: (frame_path, stats, finished)
        """
        frame_files = sorted(os.listdir("./static/result"))  # 获取所有帧文件，按顺序排列
        finished = len(frame_files) == 0
        if not frame_files:
            return None, {}, True  # 所有帧都已处理

        next_frame = frame_files.pop(0)  # 获取下一帧文件名
        frame_path = os.path.join("static/result/", next_frame)
        if frame_path != 'frame_0000.jpg':
            try:
                os.remove(frame_path)
            except PermissionError:
                time.sleep(1)

        if frame_files:
            next_frame = frame_files.pop(0)
            frame_path = os.path.join("static/result/", next_frame)
        else:
            return None, {}, True

        return frame_path, finished